# Kestrel: An XMPP-based Many-Task Computing Scheduler
# Author: Lance Stout <lancestout@gmail.com>
#
# Copyright 2009-2010 Clemson University
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""
Worker Pool Module

The worker pool module tracks the state of worker agents available to
a manager agent, such as availability and the current task executing.
"""

from __future__ import with_statement
import threading
import kestrel.core.datatypes as datatypes

class WorkerPool(object):
    """
    The worker pool module tracks the state of worker agents available to
    a manager agent, such as availability and the current task executing.

    Events listened for by this module:
      profile/approved
      presence/available
      presence/busy
      presence/offline
      job/task/start
      job/task/end

    Events raised by this module:
      task/send
      job/task/terminated
    """

    def __init__(self, kernel):
        """
        Initialize the worker pool tracking availability and
        usage data.
        """
        self.kernel = kernel
        self.events = {"profile/approved": self.approved,
                       "presence/available": self.available,
                       "presence/busy": self.busy,
                       "presence/offline": self.offline,
                       "job/task/start": self.task_started,
                       "job/task/end": self.task_finished,
                       "job/task/cancel": self.cancel_task,
                       "schedule/job": self.schedule_job,
                       "status/requested": self.status}
        self.kernel.register(self)
        self.pool = {}
        self.available = []
        self.lock = threading.RLock()


    def approved(self, worker):
        """
        A worker agent has been approved, so add it to the pool.
        """
        with self.lock:
            self.pool[worker.id] = worker
            # This is a hack to get around the fact that the worker
            # has already sent a presence notification that has probably
            # been processed before we get to this point. Now, if
            # the worker actually started out as busy, then we might
            # send an extra task request that will get rejected.
            self.kernel.event('presence/available', {'full': worker.id,
                                                     'bare': ''})


    def available(self, data):
        """Mark an agent as available."""
        with self.lock:
            worker = data['full']
            if worker in self.pool:
                prev_status = self.pool[worker].status
                self.pool[worker].status = 'available'
                self.pool[worker].task = (None, None)

                if prev_status != 'available':
                    if worker not in self.available:
                        self.available.append(worker)
                    self.kernel.event('schedule/worker', self.pool[worker])


    def busy(self, data):
        """Mark an agent as busy."""
        with self.lock:
            worker = data['full']
            if worker in self.pool:
                if worker in self.available:
                    self.available.remove(worker)
                self.pool[worker].status = 'busy'


    def offline(self, data):
        """
        Remove an agent from the pool.
        If the agent was running a task, raise a
        task/terminated event.
        """
        with self.lock:
            worker = data['full']
            if worker in self.pool:
                state, task = self.pool[worker].task
                if task is not None:
                    self.kernel.event('job/task/terminated', task)
                del self.pool[worker]
                if worker in self.available:
                    self.available.remove(worker)


    def task_started(self, task):
        """Mark an agent as running a task."""
        with self.lock:
            if task.source in self.pool:
                self.pool[task.source].task = ('running', task)


    def task_finished(self, task):
        """
        Mark an agent as not running a task because
        it just completed one.
        """
        with self.lock:
            if task.source in self.pool:
                self.pool[task.source].task = (None, None)


    def schedule_job(self, job):
        """Distribute job tasks to workers."""
        with self.lock:
            for worker in self.available:
                profile = self.pool[worker]
                if profile.matches(job.requires):
                    task = job.next_task()
                    if task is None:
                        return
                    self.pool[worker].task = ('pending', task)
                    self.kernel.event('task/send', {'to': worker,
                                                    'task': task})


    def cancel_task(self, data):
        """Notify workers to cancel a task."""
        with self.lock:
            job_id = data['job_id']
            task_id = data['task_id']

            for worker in self.pool:
                task = self.pool[worker].task[1]
                if task != None \
                        and task.job_id == job_id \
                        and task.task_id == task_id:
                    self.kernel.event('task/send/cancel', {'to':worker,
                                                           'task': task})


    def status(self, data):
        """HACK version to send status"""
        with self.lock:
            pool_size = len(self.pool)
            available = 0
            busy = 0
            for worker in self.pool:
                if self.pool[worker].status == 'available':
                    available += 1
                elif self.pool[worker].status == 'busy':
                    busy += 1

            workers = {'pool_size': pool_size,
                       'available': available,
                       'busy': busy}

            self.kernel.event('status/send', {'to': data['source'],
                                              'status': workers,
                                              'type': 'workers'})
